import logging
import numpy as np
from pinecone import Pinecone, ServerlessSpec
from tqdm import tqdm
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import time  # 导入时间库，增加等待时间

# 设置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 初始化 Pinecone
pinecone = Pinecone(api_key="18cb83ea-419c-4ac9-b60f-56155477b308")

# 索引名称
index_name = "mnist-index"

# 获取现有索引列表
existing_indexes = pinecone.list_indexes()

# 检查索引是否存在，如果存在就删除
if any(index['name'] == index_name for index in existing_indexes):
    print(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    print(f"索引 '{index_name}' 已成功删除。")
else:
    print(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
print(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
print(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
print(f"已成功连接到索引 '{index_name}'。")

# 加载 MNIST 数据
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 将数据分为训练集 (80%) 和测试集 (20%)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 上传训练数据到 Pinecone
index = pinecone.Index(index_name)
ids = [str(i) for i in range(len(X_train))]

# 使用进度条上传数据，并包含 metadata
logging.info("上传数据到 Pinecone...")
for i in tqdm(range(len(X_train)), desc="上传数据", unit="条"):
    # 上传数据时包含 metadata
    index.upsert([(ids[i], X_train[i].tolist(), {'label': int(y_train[i])})])  # 存储标签为 metadata

logging.info(f"成功上传了 {len(X_train)} 条数据。")

# 使用 k=11 测试准确率
k = 11
predictions = []

# 模拟查询 Pinecone，每个测试样本进行预测
logging.info("测试 k=11 的准确率...")
for i in tqdm(range(len(X_test)), desc="测试 k 的准确率", unit="条"):
    # 查询 Pinecone 获取前 k 个结果
    query_result = index.query(
        vector=X_test[i].tolist(),  # 测试样本向量
        top_k=k,
        include_metadata=True  # 确保包括元数据
    )

    # 从查询结果中提取预测标签
    predicted_classes = [match['metadata']['label'] for match in query_result['matches']]
    if predicted_classes:
        final_prediction = np.bincount(predicted_classes).argmax()  # 统计出现次数最多的类
        predictions.append(final_prediction)
    else:
        predictions.append(-1)  # 如果没有匹配项，可以考虑使用 -1 表示未知

# 计算准确率
accuracy = accuracy_score(y_test, predictions)
logging.info(f"当 k={k} 时的准确率为: {accuracy:.4f}")